import yaml
from sklearn.model_selection import train_test_split

import pandas as pd


def load_config(lottery):
    with open(f'{lottery}_config.yaml') as f:
        config = yaml.safe_load(f)
    return config


def preprocess_data_random_forest():
    config = load_config('ssq')
    preprocess_data_csv_path = config['data']['processed_data_path']
    data = pd.read_csv(preprocess_data_csv_path)
    # 2. 数据预处理
    X = data[['red1', 'red2', 'red3', 'red4', 'red5', 'red6']]
    y_red = data[['red1', 'red2', 'red3', 'red4', 'red5', 'red6']]  # 红球预测
    y_blue = data['blue']  # 蓝球预测

    return X, y_red, y_blue


# 3. 特征工程--随机深林
def feature_engineering_random_forest():
    X, y_red, y_blue = preprocess_data_random_forest()
    for i in range(1, 7):
        X[f'red{i}_freq'] = X[f'red{i}'].rolling(window=50).apply(
            lambda x: (x.value_counts().get(X[f'red{i}'].iloc[-1], 0)), raw=False).fillna(0)

    hot_numbers = X.apply(lambda col: col.value_counts().index[0], axis=0)
    X['hot_number_count'] = X.isin(hot_numbers).sum(axis=1)

    cold_numbers = X.apply(lambda col: col.value_counts().index[-1], axis=0)
    X['cold_number_count'] = X.isin(cold_numbers).sum(axis=1)

    # 4. 划分训练集和测试集
    X_train_red, X_test_red, y_train_red, y_test_red = train_test_split(X, y_red, test_size=0.2, random_state=42)
    X_train_blue, X_test_blue, y_train_blue, y_test_blue = train_test_split(X, y_blue, test_size=0.2, random_state=42)

    return X_train_red, X_test_red, y_train_red, y_test_red, X_train_blue, X_test_blue, y_train_blue, y_test_blue
